/*  UEFI BootLoader for EXOS
 *  Copyright (C) 2020-2022 yywd_123
 *  Author:yywd_123
 *  Date:2022-6-14
*/

#include <Uefi.h>
#include <Library/UefiLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Protocol/SimpleFileSystem.h>
#include <Guid/FileInfo.h>

#include <stdint.h>
#include <stdbool.h>

#include "BuildFlags.h"

#define Page_Size 0x1000
#define Kernel_Offset 0x100000

//  Utils.c
UINT32 strlen(const CHAR16* src);
CHAR16* strtok(CHAR16 *s, CHAR16 *delim);
UINT64 HexToDec(const CHAR16* src);

typedef struct{
  EFI_PHYSICAL_ADDRESS fbAddress;
  UINT32 *fb;

  UINT32 pitch;
  UINT32 Screen_width;
  UINT32 Screen_height;

  UINT32 BackGround_Color;
  UINT32 ForeGround_Color;
  UINT32 Cursor_x;
  UINT32 Cursor_y;
} VideoInfo;

typedef struct 
{
  unsigned long long Address;
  unsigned long long Size;
} LoadedFileInfo;

typedef struct
{
  UINT32 Type;
  UINT32 Reserved1;
  UINT64 PhysicalStart;
  UINT64 VirtualStart;
  UINT64 NumberOfPages;
  UINT64 Attribute;
  UINT64 Reserved2;
} MMAP;

typedef struct
{
  UINT64 memTotal;
} MemoryInfo;

typedef struct{
  VideoInfo Vinfo;

  LoadedFileInfo Files[6];
  UINT8 LoadedFileNum;

  MemoryInfo MemInfo;
} BootInfo;

void (*Kernel)(BootInfo *BInfo) = (void*)Kernel_Offset;

MMAP *mmap;
VOID *MMAPBuffer = NULL;
UINTN MMapSize = Page_Size * 2;
UINTN MapKey = 0;
UINTN DescSize = 0;
UINT32 DescVersion = 0;

BootInfo BInfo = {
  .Vinfo = {
    .fbAddress = 0,
    .fb = NULL,
    .pitch = 0,
    .Screen_width = 0,
    .Screen_height = 0,
    .BackGround_Color = 0xff008080,
    .ForeGround_Color = 0xffffffff,
    .Cursor_x = 0,
    .Cursor_y = 0,
  },
  .Files = {
    {0, 0},
    {0, 0},
    {0, 0},
    {0, 0},
    {0, 0},
    {0, 0},
  },
  .LoadedFileNum = 0,
  .MemInfo = {
    .memTotal = 0,
  },
};

EFI_STATUS EFIAPI ReadFile(
    IN EFI_FILE_HANDLE File,
    IN EFI_PHYSICAL_ADDRESS FileAddress)
{
  EFI_STATUS Status = 0;

  if (BInfo.LoadedFileNum >= 6)
  {
    Print(L"Too many files loaded!\n");
    return EFI_ABORTED;
  }

  EFI_FILE_INFO *FileInfo;
  UINTN BufferSize = sizeof(EFI_FILE_INFO) + sizeof(CHAR16) * 100;

  gBS->AllocatePool(
      EfiRuntimeServicesData,
      BufferSize,
      (void **)&FileInfo);

  File->GetInfo(
      File,
      &gEfiFileInfoGuid,
      &BufferSize,
      FileInfo);

  #ifdef DEBUG
  Print(L"Filename: %s\nSize: %d\nFileSize: %d\nPhysicalSize: %d\n",
      FileInfo->FileName, FileInfo->Size, FileInfo->FileSize, FileInfo->PhysicalSize);
  #endif

  UINTN PageCount = (FileInfo->FileSize + Page_Size - 1) / Page_Size;

  gBS->AllocatePages(
      AllocateAddress,
      EfiLoaderData,
      PageCount,
      &FileAddress);

  #ifdef DEBUG
  Print(L"Reading %s to Memory Address: 0x%x\n", FileInfo->FileName, FileAddress);
  #endif

  BufferSize = FileInfo->FileSize;
  Status = File->Read(
                File,
                &BufferSize,
                (void *)FileAddress);

  BInfo.Files[BInfo.LoadedFileNum].Address = FileAddress;
  BInfo.Files[BInfo.LoadedFileNum].Size = FileInfo->FileSize;
  ++BInfo.LoadedFileNum;

  gBS->FreePool(FileInfo);

  File->Close(File);

  #ifdef DEBUG
  Print(L"DataDump:\n");
  for(int i = 0; i < 50; ++i)
    Print(L"%c", *((CHAR16 *)Kernel_Offset + i));
  #endif

  return Status;
}

EFI_STATUS EFIAPI InitGraphics(
    IN EFI_HANDLE ImageHandle)
{
  EFI_STATUS Status = 0;

  EFI_GRAPHICS_OUTPUT_PROTOCOL *GOP = NULL;

  Status = gBS->LocateProtocol(
              &gEfiGraphicsOutputProtocolGuid,
              NULL,
              (void **)&GOP);
  #ifdef DEBUG
  Print(L"Init Graphics Step 1 Status: %d\n", Status);
  #endif
  if(EFI_ERROR(Status)) return Status;

  EFI_GRAPHICS_OUTPUT_MODE_INFORMATION *Info = NULL;
  UINTN InfoSize = 0;
  UINTN Screen_Size = GOP->Mode->Info->HorizontalResolution * GOP->Mode->Info->VerticalResolution;
  UINTN ModeIndex = 0;
  for(UINTN i = 0; i < GOP->Mode->MaxMode; ++i)
  {
    GOP->QueryMode(GOP, i, &InfoSize, &Info);
    if((Info->PixelFormat == 1) && (Info->HorizontalResolution * Info->VerticalResolution > Screen_Size))
    {
      Screen_Size = Info->HorizontalResolution * Info->VerticalResolution;
      ModeIndex = i;
    }
    gBS->FreePool(Info);
  }

  GOP->SetMode(GOP, ModeIndex);

  // Reload GOP Handle
  Status = gBS->LocateProtocol(
              &gEfiGraphicsOutputProtocolGuid,
              NULL,
              (void **)&GOP);
  #ifdef DEBUG
  Print(L"Init Graphics Step 2 Status: %d\n", Status);
  #endif
  if(EFI_ERROR(Status)) return Status;

  #ifdef DEBUG
  Print(L"GOP Reload Success\n");
  #endif

  BInfo.Vinfo.fbAddress = GOP->Mode->FrameBufferBase;
  BInfo.Vinfo.pitch = GOP->Mode->Info->PixelsPerScanLine;
  BInfo.Vinfo.Screen_width = GOP->Mode->Info->HorizontalResolution;
  BInfo.Vinfo.Screen_height = GOP->Mode->Info->VerticalResolution;

  return Status;
}

EFI_STATUS EFIAPI ExitBootServices(
    IN EFI_HANDLE ImageHandle)
{
  EFI_STATUS Status = 0;

  Status = gBS->AllocatePool(
              EfiLoaderData,
              MMapSize,
              &MMAPBuffer);

  Status = gBS->GetMemoryMap(
              &MMapSize,
              MMAPBuffer,
              &MapKey,
              &DescSize,
              &DescVersion);

  if(EFI_ERROR(Status)) return Status;

  for (UINTN i = (UINTN)MMAPBuffer; i < (UINTN)MMAPBuffer + MMapSize; i += DescSize)
  {
    mmap = (MMAP*)i;   
    BInfo.MemInfo.memTotal += mmap->NumberOfPages * Page_Size;
    Print(L"Type: 0x%x; Physical Start: 0x%x; Vitural Start: 0x%x; Page Count: %d; Attribute: 0x%x\n", mmap->Type, mmap->PhysicalStart, mmap->VirtualStart, mmap->NumberOfPages, mmap->Attribute);
  }

  Status = gBS->ExitBootServices(ImageHandle, MapKey);

  return Status;
}

VOID EFIAPI LoadKernel(
    IN EFI_PHYSICAL_ADDRESS Address,
    IN EFI_HANDLE ImageHandle)
{
  ExitBootServices(ImageHandle);
  
  Print(L"Memory Total: %dMB\n", BInfo.MemInfo.memTotal / 1024 / 1024);
  
  #ifdef DEBUG
  for(int i = 0; i < 50; ++i)
    Print(L"%c", *((uint8_t *)Kernel_Offset + i));
  Print(L"\n fb:0x%x", BInfo.Vinfo.fbAddress);
  #endif
  Print(L"Booting...\n");
  Kernel(&BInfo);
}

UINT32 EFIAPI GetConfigLine(
  IN OUT EFI_PHYSICAL_ADDRESS ConfigAddress,
  OUT CHAR16 *ConfigLine)
{
  UINT32 ConfigLineLength = 0;
  CHAR8 *ConfigLinePtr = (CHAR8*)ConfigAddress;
  CHAR8 c = 0;
  while((c = (*ConfigLine++ = *ConfigLinePtr++)) != L'\n' && c != L'\0')
  {
    ++ConfigLineLength;
  }
  *--ConfigLine = 0;
  ++ConfigLineLength;
  return ConfigLineLength;
}

EFI_STATUS EFIAPI UefiMain(
    IN EFI_HANDLE ImageHandle, 
    IN EFI_SYSTEM_TABLE *SystemTable)
{
  EFI_STATUS Status = EFI_SUCCESS;
  InitGraphics(ImageHandle);

  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *Volume;
  EFI_FILE_HANDLE RootFs = NULL;
  EFI_FILE_HANDLE KernelFile = NULL;
  EFI_FILE_HANDLE ConfigFile = NULL;
  
  EFI_HANDLE *HandleBuffer;
  UINTN HandleCount = 0;

  //  Get Handle
  gBS->LocateHandleBuffer(
      ByProtocol, 
      &gEfiSimpleFileSystemProtocolGuid,
      NULL,
      &HandleCount,
      &HandleBuffer);

  gBS->OpenProtocol(
      HandleBuffer[0],
      &gEfiSimpleFileSystemProtocolGuid,
      (void **)&Volume,
      ImageHandle,
      NULL,
      EFI_OPEN_PROTOCOL_GET_PROTOCOL);

  //  Read Kernel
  Volume->OpenVolume(
      Volume,
      &RootFs);

  Status = RootFs->Open(
      RootFs,
      &ConfigFile,
      L"\\EXLoader.cfg",
      EFI_FILE_MODE_READ,
      0);
  if(EFI_ERROR(Status))
  {
    Print(L"No Config File Found. Attempt to boot the kernel at \"/kernel.sys\"\n");
     Status = RootFs->Open(
                  RootFs,
                  &KernelFile,
                  L"\\kernel.sys",
                  EFI_FILE_MODE_READ,
                  0);
    if(EFI_ERROR(Status))
    {
      Print(L"No Kernel File Found.\n");
      while(1);
    }
    Print(L"Found kernel at \"/kernel.sys\".\n");
    ReadFile(KernelFile, Kernel_Offset);
    LoadKernel(Kernel_Offset, ImageHandle);

    while(1);
  }

  EFI_PHYSICAL_ADDRESS ConfigAddr = 0;
  ReadFile(ConfigFile, ConfigAddr);
  CHAR16 ConfigLine[256] = {0};
  UINT32 ConfigOffset = 0, LastOffset = 0;
  CHAR16 *ConfigArg = {0};
  BOOLEAN IsArgVaild = false;
  for(int i = 0; i < 256; ++i)
  {
    ConfigOffset += GetConfigLine(ConfigAddr + ConfigOffset, ConfigLine);
    if (ConfigOffset == LastOffset) break;
    //Print(L"%s", ConfigLine);
    ConfigArg = strtok(ConfigLine, L" ");
    while(ConfigArg != NULL) 
    {
      switch (ConfigArg[0])
      {
        case L'r':
          ConfigArg = strtok(NULL, L" ");
          CHAR16 *FilePath = ConfigArg;
          Print(L"ReadFile:%s\n", FilePath);
          ConfigArg = strtok(NULL, L" ");
          EFI_PHYSICAL_ADDRESS FileAddr = HexToDec(ConfigArg);
          Print(L"To Memory Address:0x%x\n", FileAddr);
          EFI_FILE_HANDLE File = NULL;
          Status = RootFs->Open(
                      RootFs,
                      &File,
                      FilePath,
                      EFI_FILE_MODE_READ,
                      0);
          if(EFI_ERROR(Status)) 
          {
            Print(L"File Open Error, Status:%d\nStopping execute!", Status);
            while(1);
          }
          Status = ReadFile(File, FileAddr);
          if(EFI_ERROR(Status))
          {
            Print(L"File Read Error, Status:%d\nStopping execute!", Status);
            while(1);
          }
          Print(L"Read %s to Memory Address 0x%x Success\n", FilePath, FileAddr);
          IsArgVaild = true;
          break;
        case L'l':
          EFI_PHYSICAL_ADDRESS LoadAddr = HexToDec(strtok(NULL, L" "));
          Print(L"Load Kernel at 0x%x\n", LoadAddr);
          LoadKernel(LoadAddr, ImageHandle);
          break;
        default:
          Print(L"Unknown Config Argument: %s\n", ConfigArg);
          break;
      }
      if (IsArgVaild) break;
   }
    LastOffset = ConfigOffset;
  }

  Print(L"Load Failed\n");

  while(1);

  //  Never return
  return EFI_SUCCESS;
}
